import 'package:math_keyboard/src/foundation/node.dart';
import 'package:math_keyboard/src/foundation/extended_nodes.dart';
import 'package:petitparser/petitparser.dart';

/// Extended parser for converting TeX strings to TeXNode structures.
/// 
/// This parser handles the new mathematical structures like matrices,
/// subscripts/superscripts, and limit expressions.
class ExtendedTeXParser {
  /// Constructs an [ExtendedTeXParser].
  ExtendedTeXParser(this.inputString);

  /// The TeX string to parse.
  final String inputString;

  /// Parses the TeX string and returns a TeXNode structure.
  TeXNode parse() {
    final parser = _buildParser();
    final result = parser.parse(inputString);
    
    if (result is Success) {
      return result.value;
    } else {
      throw FormatException('Failed to parse TeX: ${result.message}');
    }
  }

  /// Builds the complete parser.
  Parser<TeXNode> _buildParser() {
    return _expression().end();
  }

  /// Parser for mathematical expressions.
  Parser<TeXNode> _expression() {
    return _term().separatedBy(_operator(), includeSeparators: false)
        .map((terms) {
      if (terms.length == 1) {
        return terms.first as TeXNode;
      }
      
      // Combine terms into a single node
      final result = TeXNode(null);
      for (final term in terms) {
        if (term is TeXNode) {
          result.children.addAll(term.children);
        }
      }
      return result;
    });
  }

  /// Parser for individual terms.
  Parser<TeXNode> _term() {
    return (_matrix().cast<TeXNode>() |
           _limitExpression().cast<TeXNode>() |
           // _subscriptSuperscript().cast<TeXNode>() |
           _function().cast<TeXNode>() |
           _group() |
           _symbol() |
           _number()).cast<TeXNode>();
  }

  /// Parser for matrix structures.
  Parser<TeXMatrix> _matrix() {
    return (string(r'\begin{') &
            _matrixType() &
            string('}') &
            _matrixContent() &
            string(r'\end{') &
            _matrixType() &
            string('}'))
        .map((parts) {
      final matrixType = parts[1] as String;
      final content = parts[3] as List<List<TeXNode>>;
      
      final matrix = TeXMatrix(
        TeXNode(null),
        content.length,
        content.isNotEmpty ? content.first.length : 0,
        matrixType: matrixType,
      );
      
      // Fill matrix cells
      for (var row = 0; row < content.length; row++) {
        for (var col = 0; col < content[row].length; col++) {
          matrix.setCell(row, col, content[row][col]);
        }
      }
      
      return matrix;
    });
  }

  /// Parser for matrix types.
  Parser<String> _matrixType() {
    return (string('pmatrix') |
           string('bmatrix') |
           string('vmatrix') |
           string('Vmatrix') |
           string('matrix')).cast<String>();
  }

  /// Parser for matrix content.
  Parser<List<List<TeXNode>>> _matrixContent() {
    return _matrixRow().separatedBy(string(r' \\ '), includeSeparators: false);
  }

  /// Parser for matrix rows.
  Parser<List<TeXNode>> _matrixRow() {
    return _expression().separatedBy(string(' & '), includeSeparators: false);
  }

  /// Parser for limit expressions (integrals, sums, etc.).
  Parser<TeXLimitExpression> _limitExpression() {
    return (_limitSymbol() & _limits().optional() & _body().optional())
        .map((parts) {
      final symbol = parts[0] as String;
      final limits = parts[1] as Map<String, TeXNode>?;
      final body = parts[2] as TeXNode?;
      
      return TeXLimitExpression(
        symbol,
        TeXNode(null),
        lowerLimit: limits?['lower'],
        upperLimit: limits?['upper'],
        bodyNode: body,
        hasBody: body != null,
      );
    });
  }

  /// Parser for limit symbols.
  Parser<String> _limitSymbol() {
    return (string(r'\int') |
           string(r'\sum') |
           string(r'\prod') |
           string(r'\lim') |
           string(r'\iint') |
           string(r'\iiint') |
           string(r'\oint')).cast<String>();
  }

  /// Parser for limits (subscript and superscript).
  Parser<Map<String, TeXNode>> _limits() {
    return ((char('_') & _group() & char('^') & _group()).map((parts) => {
      'lower': parts[1] as TeXNode,
      'upper': parts[3] as TeXNode,
    }) |
    (char('^') & _group() & char('_') & _group()).map((parts) => {
      'upper': parts[1] as TeXNode,
      'lower': parts[3] as TeXNode,
    }) |
    (char('_') & _group()).map((parts) => {
      'lower': parts[1] as TeXNode,
    }) |
    (char('^') & _group()).map((parts) => {
      'upper': parts[1] as TeXNode,
    })).cast<Map<String, TeXNode>>();
  }

  /// Parser for function body.
  Parser<TeXNode> _body() {
    return _expression();
  }

  /// Parser for subscript/superscript combinations.
  // Parser<TeXSubscriptSuperscript> _subscriptSuperscript() {
  //   return (_baseExpression() & _limits())
  //       .map((parts) {
  //     final base = parts[0] as TeXNode;
  //     final limits = parts[1] as Map<String, TeXNode>;
  //
  //     return TeXSubscriptSuperscript(
  //       TeXNode(null),
  //       base,
  //       subscriptNode: limits['lower'],
  //       superscriptNode: limits['upper'],
  //     );
  //   });
  // }

  /// Parser for base expressions (for subscript/superscript).
  Parser<TeXNode> _baseExpression() {
    return (_symbol() | _group()).cast<TeXNode>();
  }

  /// Parser for functions.
  Parser<TeXFunction> _function() {
    return (_functionName() & _functionArgs())
        .map((parts) {
      final name = parts[0] as String;
      final args = parts[1] as List<TeXNode>;
      
      final argTypes = _getArgTypes(name, args.length);
      return TeXFunction(name, TeXNode(null), argTypes, args);
    });
  }

  /// Parser for function names.
  Parser<String> _functionName() {
    return (string(r'\frac') |
           string(r'\sqrt') |
           string(r'\sin') |
           string(r'\cos') |
           string(r'\tan') |
           string(r'\ln') |
           string(r'\log') |
           string(r'\underset') |
           string(r'\overset') |
           string(r'\text') |
           string(r'\mathbf') |
           string(r'\mathit') |
           string(r'\mathcal') |
           string(r'\mathbb')).cast<String>();
  }

  /// Parser for function arguments.
  Parser<List<TeXNode>> _functionArgs() {
    return _group().star();
  }

  /// Parser for grouped expressions.
  Parser<TeXNode> _group() {
    return ((char('{') & _expression() & char('}'))
        .map((parts) => parts[1] as TeXNode) |
           (char('(') & _expression() & char(')'))
        .map((parts) => parts[1] as TeXNode) |
           (char('[') & _expression() & char(']'))
        .map((parts) => parts[1] as TeXNode)).cast<TeXNode>();
  }

  /// Parser for symbols.
  Parser<TeXNode> _symbol() {
    return (_greekLetter() | _specialSymbol() | _variable()).cast<TeXNode>();
  }

  /// Parser for Greek letters.
  Parser<TeXNode> _greekLetter() {
    final greekLetters = [
      r'\alpha', r'\beta', r'\gamma', r'\delta', r'\epsilon',
      r'\zeta', r'\eta', r'\theta', r'\iota', r'\kappa',
      r'\lambda', r'\mu', r'\nu', r'\xi', r'\omicron',
      r'\pi', r'\rho', r'\sigma', r'\tau', r'\upsilon',
      r'\phi', r'\chi', r'\psi', r'\omega',
      r'\Gamma', r'\Delta', r'\Theta', r'\Lambda',
      r'\Xi', r'\Pi', r'\Sigma', r'\Upsilon',
      r'\Phi', r'\Psi', r'\Omega'
    ];
    
    var parser = string(greekLetters.first).cast<String>();
    for (var i = 1; i < greekLetters.length; i++) {
      parser = (parser | string(greekLetters[i])).cast<String>();
    }
    
    return parser.map((letter) {
      final node = TeXNode(null);
      node.addTeX(TeXLeaf(letter));
      return node;
    });
  }

  /// Parser for special symbols.
  Parser<TeXNode> _specialSymbol() {
    final specialSymbols = [
      r'\infty', r'\partial', r'\nabla', r'\therefore', r'\because',
      r'\approx', r'\neq', r'\leq', r'\geq', r'\pm', r'\mp',
      r'\leftarrow', r'\rightarrow', r'\leftrightarrow',
      r'\Leftarrow', r'\Rightarrow', r'\Leftrightarrow',
      r'\in', r'\notin', r'\subset', r'\supset', r'\cup', r'\cap'
    ];
    
    Parser<String> parser = string(specialSymbols.first).cast<String>();
    for (int i = 1; i < specialSymbols.length; i++) {
      parser = (parser | string(specialSymbols[i])).cast<String>();
    }
    
    return parser.map((symbol) {
      final node = TeXNode(null);
      node.addTeX(TeXLeaf(symbol));
      return node;
    });
  }

  /// Parser for variables.
  Parser<TeXNode> _variable() {
    return ((char('{') & letter().plus().flatten() & char('}'))
        .map((parts) {
      final variable = parts[1] as String;
      final node = TeXNode(null);
      node.addTeX(TeXLeaf('{$variable}'));
      return node;
    }) |
    letter().map((letter) {
      final node = TeXNode(null);
      node.addTeX(TeXLeaf(letter));
      return node;
    })).cast<TeXNode>();
  }

  /// Parser for numbers.
  Parser<TeXNode> _number() {
    final integer = digit().plus().flatten();
    final decimal = (integer & char('.') & integer).flatten() | integer;
    
    return decimal.map((number) {
      final node = TeXNode(null);
      for (final digit in number.split('')) {
        node.addTeX(TeXLeaf(digit));
      }
      return node;
    });
  }

  /// Parser for operators.
  Parser<String> _operator() {
    return (string(r'\times') |
           string(r'\cdot') |
           string(r'\div') |
           char('+') |
           char('-') |
           char('*') |
           char('/')).cast<String>();
  }

  /// Gets the argument types for a function.
  List<TeXArg> _getArgTypes(String functionName, int argCount) {
    switch (functionName) {
      case r'\frac':
        return [TeXArg.braces, TeXArg.braces];
      case r'\sqrt':
        return argCount == 1 
            ? [TeXArg.braces] 
            : [TeXArg.brackets, TeXArg.braces];
      case r'\underset':
      case r'\overset':
        return [TeXArg.braces, TeXArg.braces];
      case r'\text':
      case r'\mathbf':
      case r'\mathit':
      case r'\mathcal':
      case r'\mathbb':
        return [TeXArg.braces];
      default:
        return List.filled(argCount, TeXArg.braces);
    }
  }
}

/// Utility class for parsing TeX strings to TeXNode structures.
class TeXNodeParser {
  /// Parses a TeX string and returns a TeXNode.
  static TeXNode parse(String texString) {
    try {
      final parser = ExtendedTeXParser(texString);
      return parser.parse();
    } catch (e) {
      // Fallback: create a simple node with the raw TeX
      final node = TeXNode(null);
      node.addTeX(TeXLeaf(texString));
      return node;
    }
  }

  /// Parses a TeX string and returns a specific type of node.
  static T? parseAs<T extends TeX>(String texString) {
    try {
      final node = parse(texString);
      if (node.children.isNotEmpty && node.children.first is T) {
        return node.children.first as T;
      }
    } catch (e) {
      // Parsing failed
    }
    return null;
  }

  /// Checks if a TeX string represents a matrix.
  static bool isMatrix(String texString) {
    return texString.contains(r'\begin{') && 
           (texString.contains('matrix') || texString.contains('pmatrix') ||
            texString.contains('bmatrix') || texString.contains('vmatrix'));
  }

  /// Checks if a TeX string represents a limit expression.
  static bool isLimitExpression(String texString) {
    return texString.contains(r'\int') || 
           texString.contains(r'\sum') ||
           texString.contains(r'\prod') ||
           texString.contains(r'\lim');
  }

  /// Checks if a TeX string has subscripts or superscripts.
  static bool hasSubscriptSuperscript(String texString) {
    return texString.contains('_') || texString.contains('^');
  }
}